# encoding=utf-8
"""
@author: xiao nian
@contact: xiaonian030@163.com
@time: 2021-12-12 15:30
"""
from kashgari.tasks.classification import BiGRU_Model
from kashgari.utils import convert_to_saved_model
from kashgari.embeddings import BertEmbedding
from kashgari.tokenizers import BertTokenizer
import jieba
import os
from config.config import MODEL_CONFIG, TRAIN_CONFIG


def train_model(train_x,
                train_y,
                valid_x,
                valid_y,
                test_x,
                test_y):
    if MODEL_CONFIG['embedding_type'] == 'bert':
        embedding_path = os.path.abspath(MODEL_CONFIG['embedding_path'])
        embedding = BertEmbedding(embedding_path)
        tokenizer = BertTokenizer()
        tokenizer.load_from_vocab_file(os.path.join(embedding_path, 'vocab.txt'))
        train_x = [tokenizer.tokenize(s) for s in train_x]
        valid_x = [tokenizer.tokenize(s) for s in valid_x]
        test_x = [tokenizer.tokenize(s) for s in test_x]
    else:
        embedding = None
        train_x = [list(jieba.cut(s)) for s in train_x]
        valid_x = [list(jieba.cut(s)) for s in valid_x]
        test_x = [list(jieba.cut(s)) for s in test_x]

    # 初始化模型
    base_model = BiGRU_Model(embedding=embedding,
                             sequence_length=TRAIN_CONFIG['sequence_length'],
                             multi_label=MODEL_CONFIG['multi_label'])

    # 使用训练和评估数据训练模型
    # fit 方法将会返回一个 history 对象，里面有记录训练过程每一个 Epoch 的 Loss 和 Accuracy
    # 现在存储下来，用于后续的可视化
    base_history = base_model.fit(train_x,
                                  train_y,
                                  valid_x,
                                  valid_y,
                                  batch_size=TRAIN_CONFIG['batch_size'],
                                  epochs=TRAIN_CONFIG['epochs'])

    try:
        # 使用测试数据集测试模型
        # evaluate 方法输出详细的评估信息，同时以字典形式返回评估信息，存下来用于后续的比较
        base_report = base_model.evaluate(test_x, test_y)

        # 预测方法也和以前一致
        labels = base_model.predict(test_x)
    except:
        pass

    # 保存模型
    base_model.save(MODEL_CONFIG['name'])

    # 保存转换后模型到 tf_serving_model/[model_name] 目录下，版本号为 model_version
    model = BiGRU_Model.load_model(MODEL_CONFIG['name'])
    convert_to_saved_model(model, MODEL_CONFIG['directory'] + '/' + MODEL_CONFIG['name'], version=MODEL_CONFIG['version'])
